# set directory to code not slurm

library(tidyverse)
library(maptools)
library(igraph)
library(parallel)
library(redist)
library(spdep)
library(gridExtra)
library(sf)

Sys.time()
ipw <- function(x, beta, pop){
  indpop <- which(x$distance_parity <= pop)
  indbeta <- which(x$beta_sequence == beta)
  inds <- intersect(indpop, indbeta)
  psi <- x$constraint_pop[inds]
  w <- 1 / exp(beta * psi)
  inds <- sample(inds, length(inds), replace = TRUE, prob = w)
  x <- x$partitions[,inds]
  return(x)
}
ben_theme <- function(){
  theme_classic() +
    theme(panel.background = element_blank(),
          panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(),
          axis.line = element_line(colour = "black"),
          panel.border = element_rect(colour = "black", fill = NA, size = 1),
          strip.background = element_blank(),
          legend.position = "bottom", legend.title = element_blank(),
          plot.title = element_text(hjust = 0.5),
          plot.subtitle = element_text(hjust = .5))
}

fl_map <- readShapePoly("../data/fl/FL.shp")

# Load nodes
mn <- 89
nodes <- strsplit(readLines("/n/imai_lab/bfifield/enumeration/largemap_enum_all/map_sub_nodes_89.dat")," ")[[1]]

## First, subset down to right map
fl_sub <- fl_map[which(rownames(fl_map@data) %in% nodes),]

## Then, reorder nodes to be correct order to match - otherwise, won't
## line up properly with original map
fl_sub <- fl_sub[match(nodes, rownames(fl_sub@data)),]

## Convert factor columns to characters, because this particular data is messy
indx <- sapply(fl_sub@data, is.factor)
fl_sub@data[indx] <- lapply(
  fl_sub@data[indx], function(x) as.numeric(as.character(x))
)



pop <- fl_sub@data$pop
rep <- fl_sub@data$mccain
nsims <- 50000
nchains <- 4
dct <- nrow(fl_sub)
set.seed(38637) ## Random.org
Sys.time()


## Load RSG sims
load("../data/largemap_enum_all/simulations_rsg.RData")

rsg_unc <- do.call("cbind", lapply(sims_out, "[[", 1))
seg_unc_rsg <- redist.segcalc(rsg_unc, grouppop = rep, fullpop = pop)

rsg_05 <- do.call("cbind", lapply(sims_out, "[[", 2))
seg_05_rsg <- redist.segcalc(rsg_05, grouppop = rep, fullpop = pop)

rsg_01 <- do.call("cbind", lapply(sims_out, "[[", 3))
seg_01_rsg <- redist.segcalc(rsg_01, grouppop = rep, fullpop = pop)


Sys.time()
fl_sub <- st_as_sf(fl_sub)
# Calculate fryerholden
pop <- fl_sub$pop*t(matrix(rep(fl_sub$pop,nrow(fl_sub)),nrow(fl_sub)))
centroids <- st_geometry(st_centroid(x = fl_sub))
distances <- st_distance(centroids, centroids)^2




system.time(fh_unc_rsg <- unlist(mclapply(1:ncol(rsg_unc), function(x){
 ind <- rsg_unc[, x] == 1
 return(sum(pop[ind,ind]*distances[ind,ind]) + sum(pop[!ind,!ind]*distances[!ind,!ind]))
}, mc.cores = detectCores())))

write.csv(x = seg_unc_rsg, file = '/n/imai_lab/ckenny/enum_all/seg_unc_rsg.csv')
write.csv(x = fh_unc_rsg, file = '/n/imai_lab/ckenny/enum_all/fh_unc_rsg.csv')

Sys.time()
system.time(fh_01_rsg <- unlist(mclapply(1:ncol(rsg_01), function(x){
 ind <- rsg_01[, x] == 1
 return(sum(pop[ind,ind]*distances[ind,ind]) + sum(pop[!ind,!ind]*distances[!ind,!ind]))
}, mc.cores = detectCores())))


write.csv(x = seg_01_rsg, file = '/n/imai_lab/ckenny/enum_all/seg_01_rsg.csv')
write.csv(x = fh_01_rsg, file = '/n/imai_lab/ckenny/enum_all/fh_01_rsg.csv')

Sys.time()
system.time(fh_05_rsg<- unlist(mclapply(1:ncol(rsg_05), function(x){
 ind <- rsg_05[, x] == 1
 return(sum(pop[ind,ind]*distances[ind,ind]) + sum(pop[!ind,!ind]*distances[!ind,!ind]))
}, mc.cores = detectCores())))
Sys.time()

write.csv(x = seg_05_rsg, file = '/n/imai_lab/ckenny/enum_all/seg_05_rsg.csv')
write.csv(x = fh_05_rsg, file = '/n/imai_lab/ckenny/enum_all/fh_05_rsg.csv')

